#!/usr/bin/env python3
"""
诊断训练数据问题

检查：
1. 视频文件是否存在
2. 掩码文件是否存在
3. 标注点文件是否存在
4. 数据加载是否正常
"""

import json
import sys
from pathlib import Path

# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 直接使用print输出
def log_info(msg):
    print(msg)
def log_warning(msg):
    print(f"WARNING: {msg}")
def log_error(msg):
    print(f"ERROR: {msg}")

def check_sample(sample, data_dir, check_video=True, check_masks=True):
    """检查单个样本"""
    issues = []
    video_id = sample.get('id', 'unknown')
    
    # 检查视频文件
    if check_video and 'video_path' in sample:
        video_path = sample['video_path']
        # 处理Windows路径格式
        video_path = video_path.replace('\\', '/') if '\\' in video_path else video_path
        video_path_obj = Path(video_path)
        
        if not video_path_obj.exists():
            issues.append(f"视频文件不存在: {video_path}")
        elif not video_path_obj.is_file():
            issues.append(f"视频路径不是文件: {video_path}")
        else:
            pass  # 视频文件存在
    
    # 检查关键帧目录
    if 'frames_dir' in sample:
        frames_dir = Path(sample['frames_dir'])
        if not frames_dir.exists():
            issues.append(f"关键帧目录不存在: {frames_dir}")
        else:
            frame_files = list(frames_dir.glob('keyframe_*.png'))
            if len(frame_files) == 0:
                issues.append(f"关键帧目录中没有帧文件: {frames_dir}")
            else:
                pass  # 关键帧目录存在
    
    # 检查掩码和标注点
    if check_masks:
        masks_dir = data_dir / 'masks' / video_id
        annotations_dir = data_dir / 'annotations' / video_id
        
        if not masks_dir.exists():
            issues.append(f"掩码目录不存在: {masks_dir}")
        else:
            mask_files = list(masks_dir.glob('frame_*.png'))
            if len(mask_files) == 0:
                issues.append(f"掩码目录中没有掩码文件: {masks_dir}")
            else:
                pass  # 掩码目录存在
        
        if not annotations_dir.exists():
            issues.append(f"标注目录不存在: {annotations_dir}")
        else:
            ann_files = list(annotations_dir.glob('frame_*.json'))
            if len(ann_files) == 0:
                issues.append(f"标注目录中没有标注文件: {annotations_dir}")
            else:
                pass  # 标注目录存在
    
    return issues

def main():
    """主函数"""
    data_dir = Path("data/processed/segmentation")
    
    # 检查训练集
    train_split_file = data_dir / "train_split.json"
    if not train_split_file.exists():
        log_error(f"训练集划分文件不存在: {train_split_file}")
        return
    
    log_info(f"加载训练集划分文件: {train_split_file}")
    with open(train_split_file, 'r', encoding='utf-8') as f:
        train_data = json.load(f)
    
    samples = train_data.get('samples', [])
    log_info(f"总样本数: {len(samples)}")
    
    # 检查前10个样本
    log_info("=" * 60)
    log_info("检查前10个样本...")
    log_info("=" * 60)
    
    total_issues = 0
    valid_samples = 0
    
    for i, sample in enumerate(samples[:10]):
        video_id = sample.get('id', 'unknown')
        log_info(f"\n样本 {i+1}: {video_id}")
        issues = check_sample(sample, data_dir, check_video=True, check_masks=True)
        
        if issues:
            total_issues += len(issues)
            log_warning(f"  发现问题 ({len(issues)} 个):")
            for issue in issues:
                log_warning(f"    - {issue}")
        else:
            valid_samples += 1
            log_info(f"  ✓ 样本正常")
    
    log_info("=" * 60)
    log_info(f"检查结果: {valid_samples}/10 个样本正常, {total_issues} 个问题")
    
    # 统计所有样本的视频路径问题
    log_info("\n" + "=" * 60)
    log_info("统计所有样本的视频路径...")
    log_info("=" * 60)
    
    video_exists = 0
    video_not_exists = 0
    frames_dir_exists = 0
    
    for sample in samples:
        if 'video_path' in sample:
            video_path = sample['video_path']
            video_path = video_path.replace('\\', '/') if '\\' in video_path else video_path
            video_path_obj = Path(video_path)
            
            if video_path_obj.exists() and video_path_obj.is_file():
                video_exists += 1
            else:
                video_not_exists += 1
                # 检查是否有关键帧目录作为后备
                if 'frames_dir' in sample:
                    frames_dir = Path(sample['frames_dir'])
                    if frames_dir.exists():
                        frames_dir_exists += 1
    
    log_info(f"视频文件存在: {video_exists}/{len(samples)}")
    log_info(f"视频文件不存在: {video_not_exists}/{len(samples)}")
    log_info(f"有后备关键帧目录: {frames_dir_exists}/{video_not_exists}")
    
    # 统计掩码文件
    log_info("\n" + "=" * 60)
    log_info("统计掩码文件...")
    log_info("=" * 60)
    
    masks_dir = data_dir / 'masks'
    if masks_dir.exists():
        video_dirs = [d for d in masks_dir.iterdir() if d.is_dir()]
        total_masks = 0
        for video_dir in video_dirs:
            mask_files = list(video_dir.glob('frame_*.png'))
            total_masks += len(mask_files)
        log_info(f"掩码目录数: {len(video_dirs)}")
        log_info(f"总掩码文件数: {total_masks}")
    else:
        log_warning(f"掩码目录不存在: {masks_dir}")

if __name__ == '__main__':
    main()

